package com.apple.learn.nio.selector;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.Iterator;

public class SelectSockets {

	public static int PORT_NUMBER = 1234;
	
	/**
	 * @param args
	 * @throws IOException 
	 */
	public static void main(String[] args) throws IOException {
		new SelectSockets().go(args);
	}

	private void go(String[] args) throws IOException {
		int port = PORT_NUMBER;
		
		if(args.length > 0){
			port = Integer.parseInt(args[0]);
		}
		
		System.out.println("Listening on port " + port);
		
		//Allocate an unbound server socket channel
		ServerSocketChannel socketChannel = ServerSocketChannel.open();
		//Get the associated ServerSocket to bind it with
		ServerSocket serverSocket = socketChannel.socket(); 
		//Create a new Selector for use below
		Selector selector = Selector.open();
		
		//Set the port the server channel will listen to
		serverSocket.bind(new InetSocketAddress(port));
		
		//Set nonblocking mode for the listening socket
		socketChannel.configureBlocking(false);
		
		//Register the ServerSocketChannel with the Selector
		socketChannel.register(selector, SelectionKey.OP_ACCEPT);
		
		while(true){
			//This may block for a long time. Upon returning, the
			//selected set contains keys of the ready channels.
			int n = selector.select();
			
			if(n == 0){
				continue;
			}
			
			//Get an tterator over the set of selected keys
			Iterator it = selector.selectedKeys().iterator();
			
			//Look at each key in the selected set
			while(it.hasNext()){
				SelectionKey key = (SelectionKey) it.next();
				
				//Is a new connection coming in?
				if(key.isAcceptable()){
					ServerSocketChannel serverSocketChannel = (ServerSocketChannel) key.channel();
					SocketChannel channel = serverSocketChannel.accept();
					
					registerChannel(selector,channel,SelectionKey.OP_ACCEPT);
					sayHello(channel);
				}
				
				//Is there data to read on this channel?
				if(key.isReadable()){
					readDataFromSocket(key);
				}
				
				//Remove key from selected set; it's been handled 
				it.remove();
			}
		}
	}

	//Use the same byte buffer for all channels. A single thread is
	//servicing all the channels, so no danger of concurrent access.
	private ByteBuffer buffer = ByteBuffer.allocate(1024);
	
	/**
	 * Sample data handler method for a channel with data ready to read.
	 * 
	 * @param key
	 * @throws IOException 
	 */
	private void readDataFromSocket(SelectionKey key) throws IOException {
		SocketChannel socketChannel = (SocketChannel) key.channel();
		int count;
		
		buffer.clear();
		
		//Loop while data is available; channel is nonblocking
		while((count = socketChannel.read(buffer)) > 0){
			buffer.flip();//Make buffer readable
			
			//Send the data; don't assume it goes all at once
			while(buffer.hasRemaining()){
				socketChannel.write(buffer);
			}
			//
			buffer.clear();
		}
		
		if(count < 0){
			//Close channel on EOF, invalidates the key
			socketChannel.close();
		}
		
	}

	/**
	 * Register the given channel with the given selector for the given operations of interest
	 * @param selector
	 * @param channel
	 * @param opAccept
	 * @throws IOException 
	 */
	private void registerChannel(Selector selector, SocketChannel channel,
			int ops) throws IOException {
		if(channel == null){
			return;
		}
		
		//Set the new channel nonblocking
		channel.configureBlocking(false);
		
		//Register it with the selector
		channel.register(selector, ops);
	}
	
	private void sayHello(SocketChannel channel) throws IOException {
		buffer.clear();
		buffer.put("Hi there!\r\n".getBytes());
		buffer.flip();
		
		channel.write(buffer);
	}

}
